# video_client.py
# 融合视频客户端：读取相机，实时发送到服务器，并以不同帧率保存本地视频
import cv2
import socket
import numpy as np
import time
import os
import random

# ===== 全局配置变量 =====
# SERVER_HOST = "192.168.2.153"  # 服务器IP，根据实际情况修改
# SERVER_HOST = "192.168.43.17"
SERVER_HOST = "127.0.0.1"
SERVER_PORT = 6003
RECONNECT_INTERVAL = 2  # 重连间隔时间（秒）
video_target = 0

SAVE_DIR = r"."  # 视频保存目录
SAVE_TARGET_FPS = 10  # 保存视频的目标帧率（传输帧率不受此限制，使用相机自然帧率）
TRANSMIT_FPS = 10
VIDEO_FORMAT = "mp4"  # 视频格式
CV2_CODEC = "mp4v"  # OpenCV编码器

connect_timeout = 5

# 开关控制
ENABLE_TRANSMIT = True  # 是否启用传输到服务器
ENABLE_SAVE = False  # 是否启用本地保存视频
ENABLE_VISUAL = False

def connect_to_server():
    """
    连接到服务器，如果连接失败则返回None
    """
    print(f"Connecting to {SERVER_HOST}:{SERVER_PORT}...")
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(connect_timeout)
        sock.connect((SERVER_HOST, SERVER_PORT))
        print(f"Connected to {SERVER_HOST}:{SERVER_PORT}")
        return sock
    except Exception as e:
        print(f"Failed to connect to server: {e}")
        return None


def send_frame(sock, frame):
    """
    发送帧数据到服务器
    如果发送失败，返回False
    """
    try:
        # 压缩
        frame = cv2.resize(frame, (320, 240))
        # 编码为JPEG，使用扁平列表格式的params以避免兼容性问题
        ret, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
        if ret:
            data = buffer.tobytes()
            size = len(data)
            # 发送大小（4字节）+ 数据
            sock.sendall(size.to_bytes(4, "big"))
            sock.sendall(data)
        return True
    except Exception as e:
        print(f"Failed to send frame: {e}")
        return False


def fused_video_capture():
    """融合捕获：实时发送到服务器，并以指定帧率保存视频"""
    video_writer = None
    video_path = None
    
    print(f"视频目标：{video_target}")

    if ENABLE_SAVE:
        os.makedirs(SAVE_DIR, exist_ok=True)

        # 生成唯一文件名
        random_num = random.randint(1000, 9999)
        timestamp = time.strftime("%Y%m%d_%H_%M_%S")
        video_path = os.path.join(
            SAVE_DIR, f"log_video_{random_num}_{timestamp}.{VIDEO_FORMAT}"
        )

    # 打开摄像头
    cap = cv2.VideoCapture(video_target)
    if not cap.isOpened():
        raise ValueError("无法打开摄像头")

    # 获取摄像头分辨率
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    resolution = (width, height)
    print(f"摄像头分辨率: {resolution[0]}x{resolution[1]}")

    if ENABLE_SAVE:
        # 初始化视频写入器（仅用于保存，低帧率）
        fourcc = cv2.VideoWriter_fourcc(*CV2_CODEC)
        video_writer = cv2.VideoWriter(video_path, fourcc, SAVE_TARGET_FPS, resolution)

    # 帧率控制参数
    save_frame_interval = 1.0 / SAVE_TARGET_FPS if ENABLE_SAVE else float("inf")
    transmit_frame_interval = (
        1.0 / TRANSMIT_FPS if ENABLE_TRANSMIT and TRANSMIT_FPS > 0 else 0
    )  # 0表示无限制
    total_save_frames = 0
    total_transmit_frames = 0
    start_time = time.time()
    last_save_time = start_time
    last_transmit_time = start_time

    # 连接服务器
    sock = None
    if ENABLE_TRANSMIT:
        sock = connect_to_server()
        if not sock:
            print("Initial connection failed")

    transmit_desc = (
        f"~{TRANSMIT_FPS if TRANSMIT_FPS > 0 else '相机自然'}"
        if ENABLE_TRANSMIT
        else "禁用"
    )
    save_desc = f"{SAVE_TARGET_FPS}" if ENABLE_SAVE else "禁用"
    print(f"开始融合捕获：传输FPS {transmit_desc}，保存FPS: {save_desc}")
    print("按 'q' 键停止...")

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("Failed to read frame, attempting to continue...")
                time.sleep(0.1)
                continue

            current_time = time.time()
            elapsed = current_time - start_time

            # 发送帧控制（如果启用传输）
            transmit_now = False
            if ENABLE_TRANSMIT:
                transmit_now = True
                if TRANSMIT_FPS > 0:
                    if current_time - last_transmit_time < transmit_frame_interval:
                        transmit_now = False
                    else:
                        last_transmit_time = current_time

            # 如果需要发送
            if transmit_now:
                # 如果没有连接或者发送失败，则尝试重新连接
                if sock is None:
                    print(f"Reconnecting in {RECONNECT_INTERVAL} seconds...")
                    time.sleep(RECONNECT_INTERVAL)
                    sock = connect_to_server()
                    continue

                # 尝试发送帧
                if not send_frame(sock, frame):
                    print("Connection lost. Attempting to reconnect...")
                    if sock:
                        sock.close()
                    sock = None
                    continue
                total_transmit_frames += 1

            # 保存帧控制
            if ENABLE_SAVE and current_time - last_save_time >= save_frame_interval:
                video_writer.write(frame)
                total_save_frames += 1
                last_save_time = current_time

            # 显示
            if ENABLE_VISUAL:
                cv2.imshow("Video", frame)
                cv2.waitKey(1)
            
            # 显示实时状态
            current_save_fps = total_save_frames / elapsed if elapsed > 0 else 0
            status_msg = f"\r已运行: {elapsed:.1f}s | 发送帧: {total_transmit_frames} | 保存帧: {total_save_frames} | 保存FPS: {current_save_fps:.1f}"
            print(status_msg, end="")

            # 按键退出
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break

    except KeyboardInterrupt:
        print("\nClient stopped by user")
    except Exception as e:
        print(f"\nUnexpected error occurred: {e}")
    finally:
        # 释放资源
        cap.release()
        if ENABLE_VISUAL:
            cv2.destroyAllWindows()
        if ENABLE_SAVE and video_writer:
            video_writer.release()
        if sock:
            sock.close()

        total_time = time.time() - start_time
        final_save_fps = total_save_frames / total_time if total_time > 0 else 0
        print(
            f"\n捕获结束，总发送 {total_transmit_frames} 帧，总保存 {total_save_frames} 帧，平均保存帧率: {final_save_fps:.2f} FPS"
        )
        if ENABLE_SAVE and video_path:
            print(f"视频已保存到: {video_path}")


if __name__ == "__main__":
    fused_video_capture()